from typing import Optional, List, Dict
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy_crud_plus import CRUDPlus

from modules.system.models.role import Role
from modules.system.models.user_role import UserRole
from modules.system.schemas.role_schemas import RoleCreate, RoleUpdate


class RoleRepository(CRUDPlus[Role]):
    async def get_by_id(self, db: AsyncSession, pk: int) -> Optional[Role]:
        return await self.select_model(db, pk)

    async def get_by_name(self, db: AsyncSession, name: str) -> Optional[Role]:
        res = await db.execute(select(Role).where(Role.name == name))
        return res.scalar_one_or_none()

    async def create(self, db: AsyncSession, data: RoleCreate) -> Role:
        return await self.create_model(db,data)

    async def update(self, db: AsyncSession, pk: int, data: RoleUpdate) -> Optional[Role]:
        await self.update_model(db, pk, data)
        return await self.select_model(db, pk)

    async def delete(self, db: AsyncSession, pk: int) -> bool:
        return await self.delete_model(db, pk)
        
    async def get_role_ids_by_user_id(self, db: AsyncSession, user_id: int) -> List[int]:
        res = await db.execute(
            select(Role.id)
            .join(UserRole, Role.id == UserRole.role_id)
            .where(UserRole.user_id == user_id)
            .distinct()
        )
        return [row[0] for row in res.all()]

    async def get_roles_by_user_id(self, db: AsyncSession, user_id: int) -> List[Role]:
        res = await db.execute(
            select(Role)
            .join(UserRole, Role.id == UserRole.role_id)
            .where(UserRole.user_id == user_id)
            .distinct()
        )
        return [row[0] for row in res.all()]

    async def get_role_ids_map_by_user_ids(self, db: AsyncSession, user_ids: List[int]) -> Dict[int, List[int]]:
        if not user_ids:
            return {}
        res = await db.execute(
            select(UserRole.user_id, Role.id)
            .join(Role, UserRole.role_id == Role.id)
            .where(UserRole.user_id.in_(user_ids))
        )
        mapping: Dict[int, List[int]] = {}
        for uid, rid in res.all():
            mapping.setdefault(uid, []).append(rid)
        return mapping


role_dao = RoleRepository(Role)